# apps/orchestrator/orchestrator_agent.py
from typing import Optional, Dict, Any, List, Tuple
from datetime import datetime
import json
import threading
from loguru import logger

# 导入项目核心模块和常量
from utils.logger import Logger
from utils.exceptions import (
    RedisError, LockAcquireError, ParamError,
    AgentNotFoundError, AgentCallError
)
from utils.constants import (
    AgentType, TaskState, ErrorCode,
    AGENT_STATE_UPDATE_CHANNEL, TASK_STATE_UPDATE_CHANNEL
)
from middleware.redis_adapter import get_redis_adapter
from states.agent_state import AgentState
from states.task_state import TaskStateManager


class OrchestratorAgent:
    """编排Agent（多Agent协作核心调度器）"""

    def __init__(self, redis_client=None):
        """初始化编排器Agent"""
        self.redis_client = redis_client or get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.logger = Logger.get_logger(agent_name="orchestrator_agent")

        # 缓存：健康Agent实例（减少Redis查询）
        self.healthy_agents: Dict[str, List[Dict[str, Any]]] = {}  # {agent_type: [instance_info]}
        self.cache_lock = threading.Lock()  # 缓存更新锁（线程安全）

        # 启动Agent状态订阅（实时感知实例上下线）
        self._start_agent_state_subscribe()
        self.logger.success("编排Agent初始化完成，已启动Agent状态订阅")

    def _start_agent_state_subscribe(self) -> None:
        """启动Agent状态订阅（实时更新健康实例缓存）"""
        stop_event = threading.Event()

        def agent_state_callback(data: Dict[str, Any]) -> None:
            """Agent状态更新回调（实例注册/注销/更新时触发）"""
            agent_type = data.get("agent_type", "")
            instance_id = data.get("instance_id", "")
            action = data.get("action", "update")
            status = data.get("status", "")

            self.logger.debug(
                f"收到Agent状态更新: action={action} | type={agent_type} | instance={instance_id} | status={status}")

            # 加锁更新健康实例缓存
            with self.cache_lock:
                if action in ["register", "update"] and status == "healthy":
                    # 新增/更新健康实例
                    if agent_type not in self.healthy_agents:
                        self.healthy_agents[agent_type] = []
                    # 移除旧实例信息（避免重复）
                    self.healthy_agents[agent_type] = [
                        ins for ins in self.healthy_agents[agent_type]
                        if ins["instance_id"] != instance_id
                    ]
                    # 添加新实例信息
                    self.healthy_agents[agent_type].append({
                        "instance_id": instance_id,
                        "host_ip": data.get("host_ip", ""),
                        "port": data.get("port", 0),
                        "last_heartbeat": data.get("last_heartbeat", 0)
                    })
                elif action == "unregister" or status != "healthy":
                    # 移除注销/不健康实例
                    if agent_type in self.healthy_agents:
                        self.healthy_agents[agent_type] = [
                            ins for ins in self.healthy_agents[agent_type]
                            if ins["instance_id"] != instance_id
                        ]
                        # 空列表时删除键
                        if not self.healthy_agents[agent_type]:
                            del self.healthy_agents[agent_type]

        # 启动订阅线程（守护线程，随主进程退出）
        subscribe_thread = threading.Thread(
            target=self.redis_client.subscribe,
            args=(AGENT_STATE_UPDATE_CHANNEL, agent_state_callback, stop_event),
            daemon=True
        )
        subscribe_thread.start()
        self.logger.info(f"Agent状态订阅线程已启动，频道={AGENT_STATE_UPDATE_CHANNEL}")

    def _select_best_agent_instance(self, agent_type: str) -> Optional[Dict[str, Any]]:
        """选择最优Agent实例（负载均衡+健康状态优先）"""
        # 1. 先从缓存获取健康实例，缓存为空时从Redis刷新
        with self.cache_lock:
            healthy_instances = self.healthy_agents.get(agent_type, [])

        if not healthy_instances:
            self.logger.warning(f"缓存中无{agent_type}健康实例，从Redis刷新")
            healthy_instances = AgentState.get_healthy_instances(agent_type, self.redis_client)
            # 更新缓存
            with self.cache_lock:
                self.healthy_agents[agent_type] = healthy_instances

        if not healthy_instances:
            raise AgentNotFoundError(
                message=f"No healthy {agent_type} agent instance found",
                context={"agent_type": agent_type, "error_code": ErrorCode.AGENT_NOT_FOUND.value}
            )

        # 2. 负载均衡策略：选择当前任务数最少的实例
        # （从Redis获取最新的current_tasks，确保负载准确）
        best_instance = None
        min_task_count = float("inf")
        for instance in healthy_instances:
            instance_id = instance["instance_id"]
            redis_key = f"agent:state:{agent_type}:{instance_id}"
            # 从Redis获取当前任务数
            current_tasks = self.redis_client.hget(redis_key, "current_tasks")
            current_tasks = int(current_tasks) if current_tasks else 0

            # 选择任务数最少的实例
            if current_tasks < min_task_count:
                min_task_count = current_tasks
                best_instance = instance
                best_instance["current_tasks"] = current_tasks

        self.logger.debug(f"为{agent_type}选择最优实例: {best_instance['instance_id']} (当前任务数: {min_task_count})")
        return best_instance

    def _call_agent_instance(self, instance: Dict[str, Any], task_params: Dict[str, Any]) -> Dict[str, Any]:
        """调用Agent实例执行任务（模拟HTTP调用，实际可替换为真实API请求）"""
        agent_type = task_params.get("agent_type", "")
        instance_id = instance["instance_id"]
        host_ip = instance["host_ip"]
        port = instance["port"]

        self.logger.info(f"调用Agent实例: {instance_id} | 参数: {task_params}")

        try:
            # 模拟Agent执行（实际场景替换为：requests.post(f"http://{host_ip}:{port}/api/run", json=task_params)）
            # 这里返回模拟结果，与真实Agent接口返回格式一致
            simulated_result = {
                "success": True,
                "data": {
                    "task_id": task_params.get("task_id", ""),
                    "result": f"[{agent_type}] 处理成功: {task_params.get('input', '')}",
                    "processed_at": int(datetime.now().timestamp())
                },
                "error_code": 0,
                "error_msg": ""
            }

            # 模拟10%的失败概率（测试重试逻辑）
            import random
            if random.random() < 0.1:
                simulated_result["success"] = False
                simulated_result["error_code"] = ErrorCode.AGENT_CALL_ERROR.value
                simulated_result["error_msg"] = "模拟Agent调用超时"

            return simulated_result
        except Exception as e:
            self.logger.error(f"调用Agent实例{instance_id}失败: {str(e)}")
            raise AgentCallError(
                message=f"Agent instance {instance_id} call failed",
                context={"instance_id": instance_id, "error_code": ErrorCode.AGENT_CALL_ERROR.value}
            ) from e

    def _update_agent_task_metrics(self, instance: Dict[str, Any], success: bool) -> None:
        """更新Agent实例的任务指标（成功/失败次数、成功率）"""
        agent_type = instance.get("agent_type", "")
        instance_id = instance["instance_id"]
        try:
            # 初始化AgentState实例，更新指标
            agent_state = AgentState(
                agent_type=agent_type,
                instance_id=instance_id,
                redis_client=self.redis_client
            )
            agent_state.update_task_metrics(success=success)
            self.logger.debug(f"更新Agent实例{instance_id}指标成功: success={success}")
        except Exception as e:
            self.logger.warning(f"更新Agent实例{instance_id}指标失败: {str(e)}")
            # 指标更新失败不影响主流程，仅日志告警

    def _retry_task(self, task_id: str, agent_type: str, task_params: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
        """任务重试逻辑（最多重试3次）"""
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        retry_count = task_manager.get_state()["retry_count"]
        max_retries = task_manager.get_state()["max_retries"]

        if retry_count >= max_retries:
            self.logger.error(f"任务{task_id}已达到最大重试次数({max_retries})，停止重试")
            task_manager.update_state(
                state=TaskState.FAILED.value,
                error_code=ErrorCode.TASK_EXECUTE_ERROR.value,
                error_msg=f"Max retries ({max_retries}) exceeded"
            )
            return False, {"success": False, "error_msg": "Max retries exceeded"}

        # 重试计数+1，更新任务状态为重试中
        task_manager.update_state(state=TaskState.RETRYING.value)
        self.logger.info(f"任务{task_id}开始第{retry_count + 1}次重试")

        try:
            # 重新选择Agent实例（避免复用失败的实例）
            instance = self._select_best_agent_instance(agent_type)
            # 重新调用
            result = self._call_agent_instance(instance, task_params)
            if result["success"]:
                # 重试成功，更新任务状态
                task_manager.update_state(
                    state=TaskState.SUCCESS.value,
                    output_result=result["data"],
                    agent_instance_id=instance["instance_id"]
                )
                self._update_agent_task_metrics(instance, success=True)
                return True, result
            else:
                # 重试失败，继续重试或终止
                self._update_agent_task_metrics(instance, success=False)
                return self._retry_task(task_id, agent_type, task_params)
        except Exception as e:
            self.logger.error(f"任务{task_id}第{retry_count + 1}次重试失败: {str(e)}")
            return self._retry_task(task_id, agent_type, task_params)


# 全局单例实例
orchestrator_agent = OrchestratorAgent()


